import datetime as dt
import matplotlib.pyplot as plt
from matplotlib import style
from mpl_finance import candlestick_ohlc
import matplotlib.dates as mdates
import pandas as pd

plt.figure(figsize=(18, 12))

style.use('ggplot')

df = pd.read_hdf('stock_data.h5', key='data')
df = df.swaplevel('trade_date', 'ts_code')
df = df.sort_index(level=0)

df1 = df.loc['000001.SZ']
df1.index = pd.DatetimeIndex(df1.index)
df_ohlc = df1['close'].resample('5D').ohlc()

df_volume = df1['amount'].resample('5D').sum()

df_ohlc.reset_index(inplace=True)
df_ohlc['trade_date'] = df_ohlc['trade_date'].map(mdates.date2num)

ax1 = plt.subplot2grid((6,1), (0,0), rowspan=5, colspan=1)
ax2 = plt.subplot2grid((6,1), (5,0), rowspan=1, colspan=1, sharex=ax1)
ax1.xaxis_date()

candlestick_ohlc(ax1, df_ohlc.values, width=1.5, colorup='g')
ax2.fill_between(df_volume.index.map(mdates.date2num), df_volume.values, 0)

plt.savefig('./123.jpg', dpi=300, bbox_inches='tight')
plt.show()
